import os
import sys

def scan_directory(directory):
    """扫描目录，统计各类型文件的数量和占用空间"""
    ext_stats = {}
    total_count = 0
    total_size = 0

    for root, _, files in os.walk(directory):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            try:
                # 获取文件大小
                file_size = os.path.getsize(file_path)
                # 获取文件扩展名
                _, ext = os.path.splitext(file_name)
                if not ext:  # 如果没有扩展名，默认为无扩展名''
                    ext = 'no_extension'
                
                if ext in ext_stats:
                    ext_stats[ext]['count'] += 1
                    ext_stats[ext]['size'] += file_size
                else:
                    ext_stats[ext] = {'count': 1, 'size': file_size}
                
                # 更新总数和总大小
                total_count += 1
                total_size += file_size
            except Exception as e:
                print(f"无法处理文件 {file_path}: {e}")
    
    return ext_stats, total_count, total_size

def format_bytes(size):
    """格式化字节大小为可读格式"""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size < 1024:
            return f"{size:.2f} {unit}"
        size /= 1024
    return f"{size:.2f} PB"

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("使用方法: python chk_filetype.py <directory>")
        sys.exit(1)

    search_directory = sys.argv[1]
    # search_directory = "media/src"  # 替换为你的具体目录路径

    stats, total_count, total_size = scan_directory(search_directory)

    print("文件类型统计：")
    for ext, info in sorted(stats.items()):
        count = info['count']
        size = format_bytes(info['size'])
        print(f"{ext}: 数量={count}, 占用空间={size}")
    
    print(f"\n总文件数量: {total_count}, 总占用空间: {format_bytes(total_size)}")